{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "rWXrHouW3pq_"
      },
      "source": [
        "# Animation SDK example"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "cellView": "form",
        "id": "eeA1mYLdxr2j"
      },
      "outputs": [],
      "source": [
        "#@title Install the Stability SDK\n",
        "%%capture captured --no-stderr\n",
        "%pip install stability-sdk[anim]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "jCZ-IphH3prD"
      },
      "outputs": [],
      "source": [
        "#@title Mount Google Drive\n",
        "try:\n",
        "    from google.colab import drive\n",
        "    drive.mount('/content/gdrive')\n",
        "    outputs_path = \"/content/gdrive/MyDrive/AI/StableAnimation\"\n",
        "    !mkdir -p $outputs_path\n",
        "except:\n",
        "    outputs_path = \".\"\n",
        "print(f\"Animations will be saved to {outputs_path}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "zj56t6tc3prF"
      },
      "outputs": [],
      "source": [
        "#@title Connect to the Stability API\n",
        "import datetime\n",
        "import getpass\n",
        "import json\n",
        "import os\n",
        "import panel as pn\n",
        "import param\n",
        "\n",
        "from base64 import b64encode\n",
        "from IPython import display\n",
        "from tqdm import tqdm\n",
        "from types import SimpleNamespace\n",
        "\n",
        "from stability_sdk.api import Context\n",
        "from stability_sdk.animation import AnimationArgs, Animator\n",
        "from stability_sdk.utils import create_video_from_frames\n",
        "\n",
        "# @markdown To get your API key visit https://dreamstudio.ai/account\n",
        "STABILITY_HOST = \"grpc.stability.ai:443\" #@param {type:\"string\"}\n",
        "STABILITY_KEY = getpass.getpass('Enter your API Key')\n",
        "\n",
        "# Connect to Stability API\n",
        "context = Context(STABILITY_HOST, STABILITY_KEY)\n",
        "\n",
        "# Test the connection\n",
        "context.get_user_info()\n",
        "print(\"Connection successful!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ldUAFmur3prH"
      },
      "outputs": [],
      "source": [
        "# @title Settings\n",
        "\n",
        "# @markdown Run this cell to reveal the settings UI grouped across several tabs. After entering values, move on to the next step.\n",
        "\n",
        "# @markdown To reset values to default, simply re-run this cell.\n",
        "\n",
        "show_documentation = True # @param {type:'boolean'}\n",
        "\n",
        "# #@markdown ####**Resume:**\n",
        "resume_timestring = \"\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ####**Override Settings:**\n",
        "override_settings_path = \"\" #@param {type:\"string\"}\n",
        "\n",
        "###################\n",
        "\n",
        "from stability_sdk.animation import (\n",
        "    AnimationArgs,\n",
        "    Animator,\n",
        "    AnimationSettings,\n",
        "    BasicSettings,\n",
        "    CoherenceSettings,\n",
        "    ColorSettings,\n",
        "    DepthSettings,\n",
        "    InpaintingSettings,\n",
        "    Rendering3dSettings,\n",
        "    CameraSettings,\n",
        "    VideoInputSettings,\n",
        "    VideoOutputSettings,\n",
        ")\n",
        "\n",
        "args_generation = BasicSettings()\n",
        "args_animation = AnimationSettings()\n",
        "args_camera = CameraSettings()\n",
        "args_coherence = CoherenceSettings()\n",
        "args_color = ColorSettings()\n",
        "args_depth = DepthSettings()\n",
        "args_render_3d = Rendering3dSettings()\n",
        "args_inpaint = InpaintingSettings()\n",
        "args_vid_in = VideoInputSettings()\n",
        "args_vid_out = VideoOutputSettings()\n",
        "arg_objs = (\n",
        "    args_generation,\n",
        "    args_animation,\n",
        "    args_camera,\n",
        "    args_coherence,\n",
        "    args_color,\n",
        "    args_depth,\n",
        "    args_render_3d,\n",
        "    args_inpaint,\n",
        "    args_vid_in,\n",
        "    args_vid_out,\n",
        ")\n",
        "\n",
        "def _show_docs(component):\n",
        "    cols = []\n",
        "    for k, v in component.param.objects().items():\n",
        "        if k == 'name':\n",
        "            continue\n",
        "        col = pn.Column(v, v.doc)\n",
        "        cols.append(col)\n",
        "    return pn.Column(*cols)\n",
        "\n",
        "def build(component):\n",
        "    if show_documentation:\n",
        "        component = _show_docs(component)\n",
        "    return pn.Row(component, width=1000)\n",
        "\n",
        "pn.extension()\n",
        "\n",
        "pn.Tabs(*[(a.name[:-5], build(a)) for a in arg_objs])"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "_SudvbZG3prI"
      },
      "source": [
        "### Prompts"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "FT9slDSw3prJ"
      },
      "outputs": [],
      "source": [
        "animation_prompts = {\n",
        "    0: \"a painting of a delicious cheeseburger\",\n",
        "    24: \"a painting of the the answer to life the universe and everything\",\n",
        "}\n",
        "\n",
        "negative_prompt = \"\"\n",
        "negative_prompt_weight = -1.0\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Rpqv6t303prJ"
      },
      "outputs": [],
      "source": [
        "#@title Render the animation\n",
        "\n",
        "args_d = {}\n",
        "[args_d.update(a.param.values()) for a in arg_objs]\n",
        "args=AnimationArgs(**args_d)\n",
        "\n",
        "\n",
        "# load override settings if provided\n",
        "if override_settings_path:\n",
        "    if not os.path.exists(override_settings_path):\n",
        "        raise ValueError(f\"Override settings file not found: {override_settings_path}\")\n",
        "    with open(override_settings_path, 'r') as f:\n",
        "        overrides = json.load(f)\n",
        "        args = vars(args)\n",
        "        for k in args.keys():\n",
        "            if k in overrides:\n",
        "                args[k] = overrides[k]\n",
        "        args = SimpleNamespace(**args)\n",
        "        animation_prompts = overrides.get('animation_prompts', animation_prompts)\n",
        "        animation_prompts = {int(k): v for k, v in animation_prompts.items()}\n",
        "        negative_prompt = overrides.get('negative_prompt', negative_prompt)\n",
        "        negative_prompt_weight = overrides.get('negative_prompt_weight', negative_prompt_weight)\n",
        "\n",
        "# create folder for frames output\n",
        "if resume_timestring:\n",
        "    out_dir = os.path.join(outputs_path, resume_timestring)\n",
        "    if not os.path.exists(out_dir):\n",
        "        raise Exception(\"Can't resume {resume_timestring} because path {out_dir} doesn't exist. Please make sure the timestring is correct.\")\n",
        "    timestring = resume_timestring\n",
        "else:\n",
        "    timestring = datetime.datetime.now().strftime('%Y%m%d%H%M%S')\n",
        "    out_dir = os.path.join(outputs_path, timestring)\n",
        "    os.makedirs(out_dir, exist_ok=True)\n",
        "print(f\"Saving animation frames to {out_dir}...\")\n",
        "\n",
        "animator = Animator(\n",
        "    api_context=context,\n",
        "    animation_prompts=animation_prompts,\n",
        "    args=args,\n",
        "    out_dir=out_dir,    \n",
        "    negative_prompt=negative_prompt,\n",
        "    negative_prompt_weight=negative_prompt_weight,\n",
        "    resume=len(resume_timestring) != 0,\n",
        ")\n",
        "animator.save_settings(f\"{timestring}_settings.txt\")\n",
        "\n",
        "for frame in tqdm(animator.render(), initial=animator.start_frame_idx, total=args.max_frames):\n",
        "    display.clear_output(wait=True)\n",
        "    display.display(frame)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "aWhJnLNX3prL"
      },
      "outputs": [],
      "source": [
        "#@title Create video from frames\n",
        "skip_video_for_run_all = False #@param {type: 'boolean'}\n",
        "fps = 12 #@param {type:\"number\"}\n",
        "\n",
        "if skip_video_for_run_all == True:\n",
        "    print('Skipping video creation, uncheck skip_video_for_run_all if you want to run it')\n",
        "else:\n",
        "    mp4_path = os.path.join(out_dir, f\"{timestring}.mp4\")\n",
        "    print(f\"Compiling animation frames to {mp4_path}...\")\n",
        "    create_video_from_frames(out_dir, mp4_path, fps)\n",
        "\n",
        "    mp4 = open(mp4_path,'rb').read()\n",
        "    data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "    display.display( display.HTML(f'<video controls loop><source src=\"{data_url}\" type=\"video/mp4\"></video>') )"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "venv",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.5"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
